package com.yang.core.plugins.loadbalancer;

import com.yang.core.api.RestClientBuilder;
import com.yang.core.api.RestClientInterceptor;
import com.yang.core.model.RequestContext;
import com.yang.core.model.ResponseContext;
import com.yang.core.model.ExceptionContext;
import com.yang.core.extension.ExtensionPoint;
import com.yang.core.extension.RestClientExtension;
import lombok.extern.slf4j.Slf4j;

import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 负载均衡扩展
 * 支持多种负载均衡算法和健康检查
 */
@Slf4j
@ExtensionPoint(
    name = "loadBalancer",
    description = "负载均衡扩展，支持多种负载均衡算法",
    category = ExtensionPoint.Category.GENERAL,
    order = 5
)
public class LoadBalancerExtension implements RestClientExtension {

    private final LoadBalancerConfig config;
    private final LoadBalancer loadBalancer;

    public LoadBalancerExtension() {
        this.config = new LoadBalancerConfig();
        this.loadBalancer = createLoadBalancer(config);
    }

    public LoadBalancerExtension(LoadBalancerConfig config) {
        this.config = config;
        this.loadBalancer = createLoadBalancer(config);
    }

    @Override
    public void extend(RestClientBuilder builder) {
        if (config == null || !config.isEnabled()) {
            return;
        }

        builder.addInterceptor(new LoadBalancerInterceptor(loadBalancer, config));

        log.info("启用负载均衡扩展，算法: {}, 服务器数量: {}",
                config.getAlgorithm(),
                config.getServers().size());

        // 启动健康检查
        if (config.isHealthCheckEnabled()) {
            startHealthCheck();
        }
    }

    private LoadBalancer createLoadBalancer(LoadBalancerConfig config) {
        LoadBalancerStrategy strategy = createStrategy(config);
        return new LoadBalancer(strategy, config.getServers());
    }

    private LoadBalancerStrategy createStrategy(LoadBalancerConfig config) {
        return switch (config.getAlgorithm()) {
            case ROUND_ROBIN -> new RoundRobinStrategy();
            case WEIGHTED_ROUND_ROBIN -> new WeightedRoundRobinStrategy();
            case RANDOM -> new RandomStrategy();
            case LEAST_CONNECTIONS -> new LeastConnectionsStrategy();
            case IP_HASH -> new IpHashStrategy();
            case CONSISTENT_HASH -> new ConsistentHashStrategy();
            default -> throw new IllegalArgumentException("不支持的负载均衡算法: " + config.getAlgorithm());
        };
    }

    private void startHealthCheck() {
        HealthChecker healthChecker = new HealthChecker(loadBalancer, config);
        healthChecker.start();
    }

    /**
     * 负载均衡配置
     */
    public static class LoadBalancerConfig {
        private boolean enabled = true;
        private LoadBalancerAlgorithm algorithm = LoadBalancerAlgorithm.ROUND_ROBIN;
        private List<ServerInfo> servers;
        private boolean healthCheckEnabled = true;
        private int healthCheckInterval = 30000; // 30秒
        private int healthCheckTimeout = 5000;   // 5秒
        private int maxRetries = 3;
        private boolean stickySession = false;
        private String stickySessionHeader = "X-Session-ID";

        public LoadBalancerConfig() {
            this.servers = List.of();
        }

        // Getters and Setters
        public boolean isEnabled() { return enabled; }
        public void setEnabled(boolean enabled) { this.enabled = enabled; }
        public LoadBalancerAlgorithm getAlgorithm() { return algorithm; }
        public void setAlgorithm(LoadBalancerAlgorithm algorithm) { this.algorithm = algorithm; }
        public List<ServerInfo> getServers() { return servers; }
        public void setServers(List<ServerInfo> servers) { this.servers = servers; }
        public boolean isHealthCheckEnabled() { return healthCheckEnabled; }
        public void setHealthCheckEnabled(boolean healthCheckEnabled) { this.healthCheckEnabled = healthCheckEnabled; }
        public int getHealthCheckInterval() { return healthCheckInterval; }
        public void setHealthCheckInterval(int healthCheckInterval) { this.healthCheckInterval = healthCheckInterval; }
        public int getHealthCheckTimeout() { return healthCheckTimeout; }
        public void setHealthCheckTimeout(int healthCheckTimeout) { this.healthCheckTimeout = healthCheckTimeout; }
        public int getMaxRetries() { return maxRetries; }
        public void setMaxRetries(int maxRetries) { this.maxRetries = maxRetries; }
        public boolean isStickySession() { return stickySession; }
        public void setStickySession(boolean stickySession) { this.stickySession = stickySession; }
        public String getStickySessionHeader() { return stickySessionHeader; }
        public void setStickySessionHeader(String stickySessionHeader) { this.stickySessionHeader = stickySessionHeader; }

        /**
         * 添加服务器
         */
        public LoadBalancerConfig addServer(String url) {
            ServerInfo server = new ServerInfo(url);
            this.servers.add(server);
            return this;
        }

        /**
         * 添加服务器（带权重）
         */
        public LoadBalancerConfig addServer(String url, int weight) {
            ServerInfo server = new ServerInfo(url);
            server.setWeight(weight);
            this.servers.add(server);
            return this;
        }

        /**
         * 批量添加服务器
         */
        public LoadBalancerConfig addServers(List<String> urls) {
            for (String url : urls) {
                addServer(url);
            }
            return this;
        }
    }

    /**
     * 服务器信息
     */
    public static class ServerInfo {
        private String url;
        private String host;
        private int port;
        private String scheme;
        private boolean healthy = true;
        private int weight = 1;
        private AtomicInteger activeConnections = new AtomicInteger(0);
        private long lastHealthCheck;
        private Map<String, Object> metadata;

        public ServerInfo(String url) {
            this.url = url;
            this.metadata = Map.of();
            parseUrl();
        }

        private void parseUrl() {
            try {
                URI uri = URI.create(url);
                this.scheme = uri.getScheme();
                this.host = uri.getHost();
                this.port = uri.getPort() != -1 ? uri.getPort() :
                    ("https".equals(uri.getScheme()) ? 443 : 80);
            } catch (Exception e) {
                log.warn("解析服务器URL失败: {}", url, e);
                this.host = url;
                this.port = 80;
                this.scheme = "http";
            }
        }

        // Getters and Setters
        public String getUrl() { return url; }
        public String getHost() { return host; }
        public int getPort() { return port; }
        public String getScheme() { return scheme; }
        public boolean isHealthy() { return healthy; }
        public void setHealthy(boolean healthy) { this.healthy = healthy; }
        public int getWeight() { return weight; }
        public void setWeight(int weight) { this.weight = weight; }
        public AtomicInteger getActiveConnections() { return activeConnections; }
        public Map<String, Object> getMetadata() { return metadata; }
        public void setMetadata(Map<String, Object> metadata) { this.metadata = metadata; }
        public long getLastHealthCheck() { return lastHealthCheck; }
        public void setLastHealthCheck(long lastHealthCheck) { this.lastHealthCheck = lastHealthCheck; }

        public void incrementActiveConnections() {
            activeConnections.incrementAndGet();
        }

        public void decrementActiveConnections() {
            activeConnections.decrementAndGet();
        }

        @Override
        public String toString() {
            return String.format("ServerInfo{url='%s', healthy=%s, weight=%d, connections=%d}",
                    url, healthy, weight, activeConnections.get());
        }
    }

    /**
     * 负载均衡算法枚举
     */
    public enum LoadBalancerAlgorithm {
        ROUND_ROBIN,           // 轮询
        WEIGHTED_ROUND_ROBIN,  // 加权轮询
        RANDOM,                // 随机
        LEAST_CONNECTIONS,      // 最少连接
        IP_HASH,               // IP哈希
        CONSISTENT_HASH       // 一致性哈希
    }

    /**
     * 负载均衡策略接口
     */
    public interface LoadBalancerStrategy {
        ServerInfo selectServer(List<ServerInfo> servers, RequestContext context);
    }

    /**
     * 轮询策略
     */
    public static class RoundRobinStrategy implements LoadBalancerStrategy {
        private final AtomicInteger counter = new AtomicInteger(0);

        @Override
        public ServerInfo selectServer(List<ServerInfo> servers, RequestContext context) {
            if (servers == null || servers.isEmpty()) {
                return null;
            }

            // 过滤健康的服务器
            List<ServerInfo> healthyServers = servers.stream()
                .filter(ServerInfo::isHealthy)
                .toList();

            if (healthyServers.isEmpty()) {
                log.warn("没有健康的服务器可用");
                return servers.get(0); // 返回第一个服务器，即使不健康
            }

            int index = Math.abs(counter.getAndIncrement() % healthyServers.size());
            return healthyServers.get(index);
        }
    }

    /**
     * 加权轮询策略
     */
    public static class WeightedRoundRobinStrategy implements LoadBalancerStrategy {
        private final AtomicInteger counter = new AtomicInteger(0);
        private final Map<ServerInfo, AtomicInteger> currentWeights = new java.util.concurrent.ConcurrentHashMap<>();

        @Override
        public ServerInfo selectServer(List<ServerInfo> servers, RequestContext context) {
            if (servers == null || servers.isEmpty()) {
                return null;
            }

            List<ServerInfo> healthyServers = servers.stream()
                .filter(ServerInfo::isHealthy)
                .toList();

            if (healthyServers.isEmpty()) {
                log.warn("没有健康的服务器可用");
                return servers.get(0);
            }

            // 计算总权重
            int totalWeight = healthyServers.stream()
                .mapToInt(ServerInfo::getWeight)
                .sum();

            if (totalWeight <= 0) {
                return healthyServers.get(0);
            }

            // 加权轮询算法
            int index = counter.getAndIncrement() % totalWeight;
            int currentSum = 0;

            for (ServerInfo server : healthyServers) {
                currentSum += server.getWeight();
                if (index < currentSum) {
                    return server;
                }
            }

            return healthyServers.get(0);
        }
    }

    /**
     * 随机策略
     */
    public static class RandomStrategy implements LoadBalancerStrategy {
        private final java.util.Random random = new java.util.Random();

        @Override
        public ServerInfo selectServer(List<ServerInfo> servers, RequestContext context) {
            if (servers == null || servers.isEmpty()) {
                return null;
            }

            List<ServerInfo> healthyServers = servers.stream()
                .filter(ServerInfo::isHealthy)
                .toList();

            if (healthyServers.isEmpty()) {
                return servers.get(0);
            }

            int index = random.nextInt(healthyServers.size());
            return healthyServers.get(index);
        }
    }

    /**
     * 最少连接策略
     */
    public static class LeastConnectionsStrategy implements LoadBalancerStrategy {
        @Override
        public ServerInfo selectServer(List<ServerInfo> servers, RequestContext context) {
            if (servers == null || servers.isEmpty()) {
                return null;
            }

            List<ServerInfo> healthyServers = servers.stream()
                .filter(ServerInfo::isHealthy)
                .toList();

            if (healthyServers.isEmpty()) {
                return servers.get(0);
            }

            return healthyServers.stream()
                .min((s1, s2) ->
                    Integer.compare(s1.getActiveConnections().get(),
                                   s2.getActiveConnections().get()))
                .orElse(healthyServers.get(0));
        }
    }

    /**
     * IP哈希策略
     */
    public static class IpHashStrategy implements LoadBalancerStrategy {
        @Override
        public ServerInfo selectServer(List<ServerInfo> servers, RequestContext context) {
            if (servers == null || servers.isEmpty()) {
                return null;
            }

            List<ServerInfo> healthyServers = servers.stream()
                .filter(ServerInfo::isHealthy)
                .toList();

            if (healthyServers.isEmpty()) {
                return servers.get(0);
            }

            // 使用客户端IP进行哈希
            String clientIp = getClientIp(context);
            int hash = clientIp.hashCode();
            int index = Math.abs(hash) % healthyServers.size();

            return healthyServers.get(index);
        }

        private String getClientIp(RequestContext context) {
            // 从请求上下文中获取客户端IP
            String xForwardedFor = context.getHeader("X-Forwarded-For");
            if (xForwardedFor != null && !xForwardedFor.isEmpty()) {
                return xForwardedFor.split(",")[0].trim();
            }

            String xRealIp = context.getHeader("X-Real-IP");
            if (xRealIp != null && !xRealIp.isEmpty()) {
                return xRealIp;
            }

            return "unknown";
        }
    }

    /**
     * 一致性哈希策略
     */
    public static class ConsistentHashStrategy implements LoadBalancerStrategy {
        @Override
        public ServerInfo selectServer(List<ServerInfo> servers, RequestContext context) {
            if (servers == null || servers.isEmpty()) {
                return null;
            }

            List<ServerInfo> healthyServers = servers.stream()
                .filter(ServerInfo::isHealthy)
                .toList();

            if (healthyServers.isEmpty()) {
                return servers.get(0);
            }

            // 使用请求URL进行一致性哈希
            String key = context.getUrl();
            int hash = key.hashCode();
            int index = Math.abs(hash) % healthyServers.size();

            return healthyServers.get(index);
        }
    }

    /**
     * 负载均衡器
     */
    public static class LoadBalancer {
        private final LoadBalancerStrategy strategy;
        private final List<ServerInfo> servers;
        private final Map<String, ServerInfo> sessionMap = new java.util.concurrent.ConcurrentHashMap<>();

        public LoadBalancer(LoadBalancerStrategy strategy, List<ServerInfo> servers) {
            this.strategy = strategy;
            this.servers = new java.util.concurrent.CopyOnWriteArrayList<>(servers);
        }

        public ServerInfo selectServer(RequestContext context) {
            // 检查会话粘性
            String sessionId = context.getHeader("X-Session-ID");
            if (sessionId != null && !sessionId.isEmpty()) {
                ServerInfo stickyServer = sessionMap.get(sessionId);
                if (stickyServer != null && stickyServer.isHealthy()) {
                    return stickyServer;
                }
            }

            ServerInfo selected = strategy.selectServer(servers, context);

            // 记录会话粘性
            if (sessionId != null && selected != null) {
                sessionMap.put(sessionId, selected);
            }

            if (selected != null) {
                selected.incrementActiveConnections();
            }

            return selected;
        }

        public void releaseConnection(ServerInfo server) {
            if (server != null) {
                server.decrementActiveConnections();
            }
        }

        public List<ServerInfo> getServers() {
            return new java.util.ArrayList<>(servers);
        }

        public void markServerUnhealthy(ServerInfo server) {
            server.setHealthy(false);
            log.warn("标记服务器不健康: {}", server);
        }

        public void markServerHealthy(ServerInfo server) {
            server.setHealthy(true);
            server.setLastHealthCheck(System.currentTimeMillis());
            log.info("标记服务器健康: {}", server);
        }
    }

    /**
     * 健康检查器
     */
    private static class HealthChecker {
        private final LoadBalancer loadBalancer;
        private final LoadBalancerConfig config;
        private volatile boolean running = false;
        private final java.util.concurrent.ScheduledExecutorService scheduler =
            java.util.concurrent.Executors.newSingleThreadScheduledExecutor();

        public HealthChecker(LoadBalancer loadBalancer, LoadBalancerConfig config) {
            this.loadBalancer = loadBalancer;
            this.config = config;
        }

        public void start() {
            if (running) {
                return;
            }

            running = true;
            scheduler.scheduleAtFixedRate(this::performHealthCheck,
                config.getHealthCheckInterval(),
                config.getHealthCheckInterval(),
                java.util.concurrent.TimeUnit.MILLISECONDS);

            log.info("启动健康检查，间隔: {}ms", config.getHealthCheckInterval());
        }

        public void stop() {
            running = false;
            scheduler.shutdown();
            try {
                if (!scheduler.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)) {
                    scheduler.shutdownNow();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            log.info("停止健康检查");
        }

        private void performHealthCheck() {
            for (ServerInfo server : loadBalancer.getServers()) {
                checkServerHealth(server);
            }
        }

        private void checkServerHealth(ServerInfo server) {
            String healthUrl = server.getUrl() + "/health";

            try {
                // 创建一个简单的HTTP客户端进行健康检查
                java.net.HttpURLConnection connection =
                    (java.net.HttpURLConnection) new java.net.URL(healthUrl).openConnection();
                connection.setRequestMethod("GET");
                connection.setConnectTimeout(config.getHealthCheckTimeout());
                connection.setReadTimeout(config.getHealthCheckTimeout());

                int responseCode = connection.getResponseCode();
                connection.disconnect();

                if (responseCode >= 200 && responseCode < 300) {
                    if (!server.isHealthy()) {
                        loadBalancer.markServerHealthy(server);
                    }
                } else {
                    if (server.isHealthy()) {
                        loadBalancer.markServerUnhealthy(server);
                    }
                }

            } catch (Exception e) {
                if (server.isHealthy()) {
                    loadBalancer.markServerUnhealthy(server);
                }
                log.debug("健康检查失败: {} - {}", server.getUrl(), e.getMessage());
            }
        }
    }

    /**
     * 负载均衡拦截器
     */
    private static class LoadBalancerInterceptor implements RestClientInterceptor {
        private final LoadBalancer loadBalancer;
        private final LoadBalancerConfig config;

        public LoadBalancerInterceptor(LoadBalancer loadBalancer, LoadBalancerConfig config) {
            this.loadBalancer = loadBalancer;
            this.config = config;
        }

        @Override
        public void preRequest(RequestContext context) {
            // 选择服务器并修改URL
            ServerInfo selectedServer = loadBalancer.selectServer(context);
            if (selectedServer != null) {
                String originalUrl = context.getUrl();
                String newUrl = buildUrl(selectedServer, originalUrl);
                context.setUrl(newUrl);

                // 添加负载均衡信息到请求头
                context.setHeader("X-Load-Balanced-Server",
                    selectedServer.getHost() + ":" + selectedServer.getPort());

                // 保存原始服务器连接信息，用于响应后释放
                context.setAttribute("selectedServer", selectedServer);
            }
        }

        @Override
        public void postResponse(ResponseContext context) {
            // 释放服务器连接
            ServerInfo server = (ServerInfo) context.getAttribute("selectedServer");
            if (server != null) {
                loadBalancer.releaseConnection(server);
            }
        }

        @Override
        public void onError(ExceptionContext context) {
            // 释放服务器连接
            ServerInfo server = (ServerInfo) context.getAttribute("selectedServer");
            if (server != null) {
                loadBalancer.releaseConnection(server);
            }

            // 可以在这里添加故障转移逻辑
            if (config.getMaxRetries() > 0 && shouldRetry(context)) {
                log.info("尝试重试请求到其他服务器");
                // 实现重试逻辑
            }
        }

        private String buildUrl(ServerInfo server, String originalUrl) {
            if (originalUrl.startsWith("http://") || originalUrl.startsWith("https://")) {
                // 绝对URL，直接替换协议、主机和端口
                try {
                    URI uri = URI.create(originalUrl);
                    String path = uri.getPath();
                    String query = uri.getQuery();
                    StringBuilder newUrl = new StringBuilder();
                    newUrl.append(server.getScheme()).append("://");
                    newUrl.append(server.getHost()).append(":").append(server.getPort());
                    if (path != null) {
                        newUrl.append(path);
                    }
                    if (query != null) {
                        newUrl.append("?").append(query);
                    }
                    return newUrl.toString();
                } catch (Exception e) {
                    log.warn("构建URL失败，使用原始URL", e);
                    return originalUrl;
                }
            } else {
                // 相对URL，添加服务器基础URL
                String baseUrl = server.getScheme() + "://" + server.getHost() + ":" + server.getPort();
                return baseUrl.endsWith("/") ?
                    baseUrl + originalUrl.substring(1) : baseUrl + "/" + originalUrl;
            }
        }

        private boolean shouldRetry(ExceptionContext context) {
            // 实现重试判断逻辑
            return context.getException() instanceof java.net.ConnectException ||
                   context.getException() instanceof java.net.SocketTimeoutException;
        }
    }
}